import os
from flask import Flask, request, jsonify
from whoosh.index import create_in, open_dir, exists_in
from whoosh.fields import Schema, TEXT
from whoosh.qparser import QueryParser, FuzzyTermPlugin
from whoosh.analysis import StemmingAnalyzer
from openai import OpenAI
import PyPDF2
from docx import Document
import pandas as pd

# 初始化 Flask
app = Flask(__name__)

# 初始化 DeepSeek 客户端
client = OpenAI(api_key="sk-0bb2e1837ff7453c873c259368dbdf8f", base_url="https://api.deepseek.com")

# 定义知识库文件夹路径
KNOWLEDGE_BASE_FOLDER = r"E:\\BijiApp\\web_app\\python_part\\Demo1\\TEMP"

# 定义 Whoosh 索引路径
INDEX_DIR = "whoosh_index"

# 定义支持模糊匹配的分析器
analyzer = StemmingAnalyzer()

def create_search_index(folder_path, index_dir):
    """创建或更新 Whoosh 倒排索引，支持模糊匹配"""
    schema = Schema(content=TEXT(stored=True, analyzer=analyzer))
    if not os.path.exists(index_dir):
        os.mkdir(index_dir)
    index = create_in(index_dir, schema)
    writer = index.writer()

    for filename in os.listdir(folder_path):
        file_path = os.path.join(folder_path, filename)
        content = ""
        print(f"Processing file: {filename}")

        if filename.endswith('.txt'):
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
        elif filename.endswith('.pdf'):
            with open(file_path, 'rb') as file:
                reader = PyPDF2.PdfReader(file)
                for page in reader.pages:
                    content += page.extract_text()
        elif filename.endswith('.docx'):
            doc = Document(file_path)
            for paragraph in doc.paragraphs:
                content += paragraph.text + "\n"
        elif filename.endswith('.xlsx'):
            df = pd.read_excel(file_path)
            content = df.to_string(index=False)

        print(f"Content: {content[:100]}...")  # 打印前 100 个字符
        if content:
            writer.add_document(content=content)
    writer.commit()
    print("Index updated successfully.")

def search_index(query, index_dir, top_k=3):
    """在 Whoosh 索引中检索相关文档，支持模糊匹配"""
    index = open_dir(index_dir)
    with index.searcher() as searcher:
        query_parser = QueryParser("content", index.schema)
        query_parser.add_plugin(FuzzyTermPlugin())
        query = query_parser.parse(query + "~2")  # 允许最多 2 个字符的编辑距离
        print(f"Parsed Query: {query}")  # 打印解析后的查询
        results = searcher.search(query, limit=top_k)
        print(f"Search Results: {results}")  # 打印检索结果
        return [result['content'] for result in results]

def call_deepseek_api(query, context):
    """通过 OpenAI SDK 调用 DeepSeek API"""
    try:
        messages = [{"role": "user", "content": f"{context}\n\n问题：{query}"}]
        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=messages
        )
        print("Request Data:", messages)
        print("Response:", response)
        return response.choices[0].message.content
    except Exception as e:
        print("Error:", e)
        return None

# 创建 Whoosh 索引（如果索引不存在）
if not os.path.exists(INDEX_DIR):
    print("Creating search index...")
    create_search_index(KNOWLEDGE_BASE_FOLDER, INDEX_DIR)

@app.route('/query', methods=['POST'])
def query():
    """处理用户查询"""
    user_query = request.json.get('query')
    relevant_docs = search_index(user_query, INDEX_DIR, top_k=3)
    context = "\n".join(relevant_docs)
    deepseek_response = call_deepseek_api(user_query, context)
    return jsonify({
        'query': user_query,
        'context': context,
        'response': deepseek_response
    })

@app.route('/update_index', methods=['POST'])
def update_index():
    """更新 Whoosh 索引"""
    try:
        create_search_index(KNOWLEDGE_BASE_FOLDER, INDEX_DIR)
        return jsonify({"status": "success", "message": "Index updated successfully."})
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)})

if __name__ == '__main__':
    app.run(port=5000)